#!/usr/bin/env python3
"""
RocketMQ 生产级消息消费者

功能特性：
- 多线程并发消费
- 消息处理失败重试
- 优雅关闭机制
- 消费进度监控
- 死信队列处理
- 消费限流控制
- 健康检查和故障恢复

Author: Python RocketMQ Guide
Date: 2025-08-18
"""

import logging
import threading
import time
from abc import ABC, abstractmethod
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from enum import Enum
from typing import Dict, Any, Optional, List, Union, Callable, Set
from queue import Queue, Empty
import signal

try:
    from rocketmq.client import PullConsumer, ConsumeStatus, Message
    from rocketmq.client import MessageQueue, PullResult, PullStatus
except ImportError:
    raise ImportError("请安装 rocketmq-client-python: pip install rocketmq-client-python")

import weakref
import traceback
from threading import Event, RLock, Semaphore

logger = logging.getLogger(__name__)


class ConsumeResult(Enum):
    """消费结果枚举"""
    SUCCESS = "SUCCESS"
    RETRY = "RETRY"
    FAILED = "FAILED"
    SKIP = "SKIP"


@dataclass
class ConsumeContext:
    """消费上下文"""
    message: Message
    topic: str
    message_queue: MessageQueue
    queue_offset: int
    retry_count: int = 0
    max_retries: int = 3
    consume_start_time: float = None
    properties: Optional[Dict[str, Any]] = None

    def __post_init__(self):
        if self.consume_start_time is None:
            self.consume_start_time = time.time()


@dataclass
class ConsumeStats:
    """消费统计信息"""
    consumed_messages: int = 0
    failed_messages: int = 0
    retry_messages: int = 0
    skip_messages: int = 0
    total_consume_time: float = 0
    last_consume_time: Optional[float] = None

    @property
    def success_rate(self) -> float:
        """成功率"""
        total = self.consumed_messages + self.failed_messages
        return self.consumed_messages / total if total > 0 else 0

    @property
    def avg_consume_time(self) -> float:
        """平均消费时间"""
        return self.total_consume_time / self.consumed_messages if self.consumed_messages > 0 else 0


class MessageHandler(ABC):
    """消息处理器抽象基类"""

    @abstractmethod
    def handle_message(self, context: ConsumeContext) -> ConsumeResult:
        """
        处理消息

        Args:
            context: 消费上下文

        Returns:
            ConsumeResult: 处理结果
        """
        pass

    def on_consume_success(self, context: ConsumeContext):
        """消费成功回调"""
        pass

    def on_consume_failed(self, context: ConsumeContext, error: Exception):
        """消费失败回调"""
        pass


class RateLimiter:
    """速率限制器"""

    def __init__(self, max_rate: float, window_size: float = 1.0):
        """
        初始化速率限制器

        Args:
            max_rate: 最大速率（每秒）
            window_size: 时间窗口大小（秒）
        """
        self.max_rate = max_rate
        self.window_size = window_size
        self.tokens = max_rate
        self.last_update = time.time()
        self._lock = RLock()

    def acquire(self, tokens: int = 1) -> bool:
        """
        获取令牌

        Args:
            tokens: 需要的令牌数

        Returns:
            bool: 是否获取成功
        """
        with self._lock:
            now = time.time()
            # 添加新的令牌
            elapsed = now - self.last_update
            self.tokens = min(self.max_rate, self.tokens + elapsed * self.max_rate)
            self.last_update = now

            # 检查是否有足够的令牌
            if self.tokens >= tokens:
                self.tokens -= tokens
                return True
            return False

    def wait_for_token(self, tokens: int = 1, timeout: Optional[float] = None) -> bool:
        """
        等待令牌

        Args:
            tokens: 需要的令牌数
            timeout: 超时时间

        Returns:
            bool: 是否获取成功
        """
        start_time = time.time()
        while True:
            if self.acquire(tokens):
                return True

            if timeout and (time.time() - start_time) > timeout:
                return False

            time.sleep(0.01)  # 短暂等待


class MessageConsumer:
    """生产级消息消费者"""

    def __init__(self, config=None, handler: MessageHandler = None):
        """
        初始化消费者

        Args:
            config: RocketMQ 配置
            handler: 消息处理器
        """
        # 配置管理
        if config is None:
            config = self._get_default_config()
        self.config = config
        self.handler = handler

        # 消费者实例
        self.consumer = None
        self.subscribed_topics: Set[str] = set()

        # 线程管理
        self.consume_executor = ThreadPoolExecutor(
            max_workers=config.consumer_thread_count,
            thread_name_prefix="rocketmq-consumer"
        )

        # 状态管理
        self._running = False
        self._shutdown_event = Event()
        self._consume_threads = []
        self._stats_lock = RLock()

        # 统计信息
        self._stats = ConsumeStats()

        # 速率限制
        self.rate_limiter = None
        if hasattr(config, 'consume_rate_limit') and config.consume_rate_limit > 0:
            self.rate_limiter = RateLimiter(config.consume_rate_limit)

        # 重试队列
        self._retry_queue = Queue()
        self._dead_letter_queue = Queue()

        # 信号处理
        self._setup_signal_handlers()

        logger.info("MessageConsumer 初始化完成")

    def _get_default_config(self):
        """获取默认配置"""
        from types import SimpleNamespace

        config = SimpleNamespace()
        config.name_server_address = "localhost:9876"
        config.consumer_group = "DefaultConsumerGroup"
        config.consumer_message_model = "CLUSTERING"
        config.consumer_thread_count = 20
        config.consumer_max_reconsume_times = 16
        config.consumer_pull_timeout = 30000
        config.consumer_pull_batch_size = 32
        config.access_key = None
        config.secret_key = None
        config.enable_health_check = True
        config.health_check_interval = 30

        return config

    def _setup_signal_handlers(self):
        """设置信号处理器"""

        def signal_handler(signum, frame):
            logger.info(f"接收到信号 {signum}，开始优雅关闭")
            self.shutdown()

        signal.signal(signal.SIGTERM, signal_handler)
        signal.signal(signal.SIGINT, signal_handler)

    def _initialize_consumer(self):
        """初始化消费者实例"""
        if self.consumer is not None:
            return

        try:
            self.consumer = PullConsumer(self.config.consumer_group)
            self.consumer.set_name_server_address(self.config.name_server_address)

            # 设置消费模式
            if self.config.consumer_message_model == "BROADCASTING":
                self.consumer.set_message_model("BROADCASTING")
            else:
                self.consumer.set_message_model("CLUSTERING")

            # 设置认证信息
            if self.config.access_key and self.config.secret_key:
                self.consumer.set_session_credentials(
                    self.config.access_key,
                    self.config.secret_key,
                    ""
                )

            self.consumer.start()
            logger.info("消费者实例初始化成功")

        except Exception as e:
            logger.error(f"初始化消费者失败: {e}")
            raise

    def subscribe(self, topic: str, tags: str = "*"):
        """
        订阅主题

        Args:
            topic: 主题名称
            tags: 标签过滤器
        """
        if not self._running:
            self._initialize_consumer()

        try:
            # 注意：PullConsumer 不直接支持 subscribe，需要手动获取队列
            self.subscribed_topics.add(topic)
            logger.info(f"订阅主题: {topic}, 标签: {tags}")

        except Exception as e:
            logger.error(f"订阅主题失败 {topic}: {e}")
            raise

    def start_consuming(self):
        """开始消费消息"""
        if self._running:
            logger.warning("消费者已在运行中")
            return

        if not self.subscribed_topics:
            raise ValueError("请先订阅至少一个主题")

        if not self.handler:
            raise ValueError("请设置消息处理器")

        self._running = True
        self._shutdown_event.clear()

        # 启动消费线程
        for topic in self.subscribed_topics:
            thread = threading.Thread(
                target=self._consume_topic,
                args=(topic,),
                daemon=True,
                name=f"consumer-{topic}"
            )
            thread.start()
            self._consume_threads.append(thread)

        # 启动重试处理线程
        retry_thread = threading.Thread(
            target=self._process_retry_messages,
            daemon=True,
            name="retry-processor"
        )
        retry_thread.start()
        self._consume_threads.append(retry_thread)

        # 启动健康检查
        if self.config.enable_health_check:
            health_thread = threading.Thread(
                target=self._health_check_loop,
                daemon=True,
                name="health-checker"
            )
            health_thread.start()
            self._consume_threads.append(health_thread)

        logger.info("消息消费已启动")

    def _consume_topic(self, topic: str):
        """消费指定主题的消息"""
        logger.info(f"开始消费主题: {topic}")

        while not self._shutdown_event.is_set():
            try:
                # 获取主题的消息队列
                message_queues = self.consumer.fetch_subscribe_message_queues(topic)

                if not message_queues:
                    logger.warning(f"主题 {topic} 没有可用的消息队列")
                    time.sleep(5)
                    continue

                # 并发处理多个队列
                futures = []
                for mq in message_queues:
                    future = self.consume_executor.submit(self._consume_queue, topic, mq)
                    futures.append(future)

                # 等待一轮消费完成
                for future in as_completed(futures, timeout=10):
                    try:
                        future.result()
                    except Exception as e:
                        logger.error(f"消费队列失败: {e}")

            except Exception as e:
                logger.error(f"消费主题 {topic} 出错: {e}")
                if not self._shutdown_event.is_set():
                    time.sleep(5)  # 出错后等待5秒再重试

    def _consume_queue(self, topic: str, message_queue: MessageQueue):
        """消费指定队列的消息"""
        try:
            # 拉取消息
            pull_result = self.consumer.pull_block_if_not_found(
                message_queue,
                None,  # 使用默认的队列偏移量
                self.config.consumer_pull_batch_size,
                self.config.consumer_pull_timeout
            )

            if pull_result.pull_status != PullStatus.FOUND:
                return

            # 处理拉取到的消息
            for message in pull_result.msg_found_list:
                if self._shutdown_event.is_set():
                    break

                # 速率限制
                if self.rate_limiter and not self.rate_limiter.wait_for_token(timeout=1):
                    logger.warning("达到速率限制，跳过消息处理")
                    continue

                # 创建消费上下文
                context = ConsumeContext(
                    message=message,
                    topic=topic,
                    message_queue=message_queue,
                    queue_offset=message.queue_offset
                )

                # 提交到线程池处理
                self.consume_executor.submit(self._handle_message, context)

        except Exception as e:
            if not self._shutdown_event.is_set():
                logger.error(f"拉取消息失败: {e}")

    def _handle_message(self, context: ConsumeContext):
        """处理单条消息"""
        start_time = time.time()

        try:
            # 调用业务处理器
            result = self.handler.handle_message(context)

            # 处理结果
            elapsed_time = time.time() - start_time

            if result == ConsumeResult.SUCCESS:
                self._handle_consume_success(context, elapsed_time)
                self.handler.on_consume_success(context)

            elif result == ConsumeResult.RETRY:
                self._handle_consume_retry(context)

            elif result == ConsumeResult.SKIP:
                self._handle_consume_skip(context, elapsed_time)

            else:  # FAILED
                self._handle_consume_failed(context, elapsed_time, None)

        except Exception as e:
            elapsed_time = time.time() - start_time
            logger.error(f"处理消息异常: {e}\n{traceback.format_exc()}")
            self._handle_consume_failed(context, elapsed_time, e)
            self.handler.on_consume_failed(context, e)

    def _handle_consume_success(self, context: ConsumeContext, elapsed_time: float):
        """处理消费成功"""
        with self._stats_lock:
            self._stats.consumed_messages += 1
            self._stats.total_consume_time += elapsed_time
            self._stats.last_consume_time = time.time()

        logger.debug(f"消息消费成功: {context.message.msg_id}, 耗时: {elapsed_time:.3f}s")

    def _handle_consume_retry(self, context: ConsumeContext):
        """处理消费重试"""
        if context.retry_count < context.max_retries:
            context.retry_count += 1
            self._retry_queue.put(context)

            with self._stats_lock:
                self._stats.retry_messages += 1

            logger.info(f"消息进入重试队列: {context.message.msg_id}, 重试次数: {context.retry_count}")
        else:
            # 超过最大重试次数，进入死信队列
            self._dead_letter_queue.put(context)
            logger.warning(f"消息超过最大重试次数，进入死信队列: {context.message.msg_id}")

    def _handle_consume_skip(self, context: ConsumeContext, elapsed_time: float):
        """处理消费跳过"""
        with self._stats_lock:
            self._stats.skip_messages += 1
            self._stats.total_consume_time += elapsed_time

        logger.info(f"消息被跳过: {context.message.msg_id}")

    def _handle_consume_failed(self, context: ConsumeContext, elapsed_time: float, error: Optional[Exception]):
        """处理消费失败"""
        with self._stats_lock:
            self._stats.failed_messages += 1
            self._stats.total_consume_time += elapsed_time

        logger.error(f"消息消费失败: {context.message.msg_id}, 错误: {error}")

        # 根据配置决定是否重试
        if context.retry_count < context.max_retries:
            self._handle_consume_retry(context)

    def _process_retry_messages(self):
        """处理重试消息"""
        logger.info("重试消息处理线程已启动")

        while not self._shutdown_event.is_set():
            try:
                # 从重试队列获取消息
                try:
                    context = self._retry_queue.get(timeout=1)
                except Empty:
                    continue

                # 计算重试延迟
                delay = min(2 ** context.retry_count, 60)  # 最大延迟60秒

                logger.info(f"重试消息 {context.message.msg_id}，延迟 {delay} 秒")
                time.sleep(delay)

                # 重新处理消息
                if not self._shutdown_event.is_set():
                    self._handle_message(context)

            except Exception as e:
                logger.error(f"处理重试消息失败: {e}")

    def _health_check_loop(self):
        """健康检查循环"""
        logger.info("健康检查线程已启动")

        while not self._shutdown_event.wait(self.config.health_check_interval):
            try:
                self._perform_health_check()
            except Exception as e:
                logger.error(f"健康检查失败: {e}")

    def _perform_health_check(self):
        """执行健康检查"""
        # 检查消费统计
        stats = self.get_stats()

        # 检查最近是否有消费活动
        if stats['last_consume_time']:
            last_consume_elapsed = time.time() - stats['last_consume_time']
            if last_consume_elapsed > 300:  # 5分钟内没有消费活动
                logger.warning(f"长时间未消费消息: {last_consume_elapsed:.1f} 秒")

        # 检查重试队列大小
        retry_queue_size = self._retry_queue.qsize()
        if retry_queue_size > 100:
            logger.warning(f"重试队列积压严重: {retry_queue_size} 条消息")

        # 检查死信队列大小
        dead_letter_queue_size = self._dead_letter_queue.qsize()
        if dead_letter_queue_size > 0:
            logger.warning(f"死信队列有消息: {dead_letter_queue_size} 条")

        logger.debug("健康检查完成")

    def get_stats(self) -> Dict[str, Any]:
        """获取消费统计信息"""
        with self._stats_lock:
            stats = {
                "consumed_messages": self._stats.consumed_messages,
                "failed_messages": self._stats.failed_messages,
                "retry_messages": self._stats.retry_messages,
                "skip_messages": self._stats.skip_messages,
                "success_rate": self._stats.success_rate,
                "avg_consume_time": self._stats.avg_consume_time,
                "last_consume_time": self._stats.last_consume_time,
                "retry_queue_size": self._retry_queue.qsize(),
                "dead_letter_queue_size": self._dead_letter_queue.qsize(),
                "running": self._running,
                "subscribed_topics": list(self.subscribed_topics)
            }

        return stats

    def shutdown(self, timeout: float = 30):
        """
        优雅关闭消费者

        Args:
            timeout: 关闭超时时间（秒）
        """
        if not self._running:
            return

        logger.info("开始关闭消费者...")

        # 设置关闭标志
        self._running = False
        self._shutdown_event.set()

        # 等待所有消费线程结束
        for thread in self._consume_threads:
            thread.join(timeout=timeout / len(self._consume_threads) if self._consume_threads else timeout)
            if thread.is_alive():
                logger.warning(f"线程 {thread.name} 未能在超时时间内结束")

        # 关闭线程池
        self.consume_executor.shutdown(wait=True)

        # 关闭消费者实例
        if self.consumer:
            try:
                self.consumer.shutdown()
            except Exception as e:
                logger.error(f"关闭消费者实例失败: {e}")

        logger.info("消费者已关闭")

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.shutdown()


class DefaultMessageHandler(MessageHandler):
    """默认消息处理器实现"""

    def __init__(self, message_processor: Callable[[ConsumeContext], ConsumeResult]):
        """
        初始化默认处理器

        Args:
            message_processor: 消息处理函数
        """
        self.message_processor = message_processor

    def handle_message(self, context: ConsumeContext) -> ConsumeResult:
        """处理消息"""
        try:
            return self.message_processor(context)
        except Exception as e:
            logger.error(f"消息处理异常: {e}")
            return ConsumeResult.RETRY
